# -*- coding: utf-8 -*-
import os
import math
from jsonpath import jsonpath
from locust import LoadTestShape
from locust import HttpUser, task, constant


class QuickStartUser(HttpUser):
    host = "http://127.0.0.1:8000"
    wait_time = constant(1)

    @task
    def api(self):
        with self.client.get(url="/locust_test", catch_response=True, name='测试') as response:
            if response.status_code == 200:
                try:
                    assert jsonpath(response.json(), '$..code')[0] == 1
                except:
                    response.failure("Timeout")
            else:
                response.failure(response.status_code)


class StagesShape(LoadTestShape):
    """
    自定义一个双波形图形，
    模拟在某两个时间点的最高值


    参数解析:
        min_users ： 最小用户数
        peak_one_users ： 用户在第一个峰值
        peak_two_users ： 用户在第二个峰值
        time_limit ： 测试执行总时间


    """
    # 最小用户数
    min_users = 20

    # 第一个峰值的用户数
    peak_one_users = 60

    # 第二个峰值的用户数
    peak_two_users = 40

    # 测试执行时间
    time_limit = 600

    def tick(self):

        # 将get_run_time 四舍五入
        run_time = round(self.get_run_time())

        if run_time < self.time_limit:

            user_count = (
                    (self.peak_one_users - self.min_users)
                    * math.e ** -(((run_time / (self.time_limit / 10 * 2 / 3)) - 5) ** 2)
                    + (self.peak_two_users - self.min_users)
                    * math.e ** -(((run_time / (self.time_limit / 10 * 2 / 3)) - 10) ** 2)
                    + self.min_users

            )
            tick_data = (round(user_count), round(user_count))
            return tick_data
        else:
            return None


if __name__ == "__main__":
    os.system("locust -f 负载图_双波形策略.py --web-host=127.0.0.1")
